from paddlex import create_pipeline

pipeline = create_pipeline(
    pipeline="PP-ChatOCRv3-doc",
    llm_name="ernie-4.0",
    llm_params={"api_type": "aistudio", "access_token": "9b0c7f625a8badfbc19691472e336ec582ea04a4"},
    layout_model="./official_models/PicoDet-L_layout_3cls",  # 版面区域检测模块模型s
    table_model="./official_models/SLANet_plus",  # 表格识别模型
    # text_det_model="./official_models/PP-OCRv4_server_det",  # 文本检测模块模型 PP-OCRv4_server_det
    # text_rec_model="./official_models/PP-OCRv4_server_rec",  # 文本识别模块模型 PP-OCRv4_server_rec
    text_det_model="./official_models/PP-OCRv4_mobile_det",  # 文本检测模块模型 PP-OCRv4_server_det
    text_rec_model="./official_models/PP-OCRv4_mobile_rec",  # 文本识别模块模型 PP-OCRv4_server_rec
    seal_text_det_model="./official_models/PP-OCRv4_mobile_seal_det",  # 印章文本检测
    # doc_image_ori_cls_model="PP-LCNet_x1_0_doc_ori",  # 文档图像方向分类模块模型
    text_det_batch_size=1,  # 该参数只能是1
    text_rec_batch_size=5
)

print("------1--1660838898407034880.jpg---")
# visual_result, visual_info = pipeline.visual_predict("./cards/bank/1660838898407034880.jpg")
# visual_result, visual_info = pipeline.visual_predict("./cards/id/1338402204291276800.png")
# visual_result, visual_info = pipeline.visual_predict("./cards/license/1338330560398237696.png")
# visual_result, visual_info = pipeline.visual_predict("./cards/weightTicket/1.58.jpg")
visual_result, visual_info = pipeline.visual_predict("./image/27.png")
print(visual_info)

print("------1--1660838898407034880.jpg---")
print(visual_result)

for res in visual_result:
    res.save_to_img("./output")
    res.save_to_html('./output')
    res.save_to_xlsx('./output')

vector = pipeline.build_vector(visual_info=visual_info)

text_block = visual_info['ocr_text'][0]['words in text block']

kIdCard = [
    "请给出正确的身份证信息,包含:姓名、性别、民族、出生、身份证号、住址，身份证号必须是完整的18位数字,格式为 {type:ID,name:xxx,sex:xxx,nation:xxx ,birth:xxx,idNo:xxxxxxxxxxxxxxxxx ,address:xxx}"]

# kBankCard = [
#     "请给出正确的银行卡号和银行名称信息,银行卡号必须是完整的19位数字或者16位数字,"
#     "由前面6位数字加后面13位数字构成的19位数字或者由连续的4位构成的16位数字并且连续的4位数字不会重复，"
#     "银行名称必须是中文，格式为 {type:bank,cardNo:xxxxxx xxxxxxxxxxxxx or xxxx xxxx xxxx xxxx，cardName:招商银行} carNO 即银行卡号 必须是19位或者16位数字"]
kBankCard = ["给出银行卡卡片信息，格式为  res:{xx:xx}"]
kLicense = [
    "请给出正确的营业执照信息,包含:营业执照注册号即（统一社会信用代码）、公司名称即（名称）、公司类型、法定代表人、经营范围、注册资本、成立日期、营业期限、住所即（公司地址），营业执照注册号必须是完整的18位数字,"
    "格式为 {type:license,licenseNo:xxxxxxxxxxxxxxxxx,companyName:xxx,companyType:xxxx,corporation:xxx,scope:xxx,registeredCapital:xxx,establishmentDate:xxx,businessPeriod:xxx,companyAddress:xxx}"]

# kWeightTicket = [
#     "从图片中提取所有文本信息，"
#     "表格第一行包括、供应商、车号、日期、皮重过磅时间、皮重过磅重量、皮重过磅员、"
#     "表格第二行毛重过磅时间、毛重过磅重量、毛重过磅员、净重、"
#     "表格第三行货物信息 包括 货物名称、结算重量等。"
#     "以结构化的格式返回数据，"
#     "格式为 {供应商:xxx,车号:xxx,日期:xxx,皮重过磅时间:xxx,皮重过磅重量:xxx,皮重过磅员:xxx,毛重过磅时间:xxx,毛重过磅重量:xxx,毛重过磅员:xxx,净重:xxx,货物名称:xxx,"
#     "结算重量:xxx}"
#     ""]
kWeightTicket = ["从图片中提取所有文本信息,输出markdown格式的表格数据"]
## 包含的关键词

weight_ticket_key_word = ["过磅单"]

bank_key_word = ["Debit", "DEBIT", "ATM", "银联", "BOC", "CCB", "ABC", "ICBC", "BOCOM", "CMB", "CIB", "CMBC", "HXB",
                 "PAB",
                 "CEB", "CNCB", "SPDB"]

id_key_word = ["公民身份号码"]

license_key_word = ["法定代表人"]

bank_advice = ["回单"]

keyword = []
# 判断文本中是否包含关键词
if all(key in text_block for key in license_key_word):
    keyword = kLicense
elif all(key in text_block for key in id_key_word):
    keyword = kIdCard
elif all(key in text_block for key in weight_ticket_key_word):
    keyword = kWeightTicket
elif all(key in text_block for key in bank_advice):
    keyword = ""
else:
    keyword = kBankCard

print(keyword)
print(text_block)
chat_result = pipeline.chat(
    key_list=[] + keyword,
    visual_info=visual_info,
    vector=vector,
)

print(chat_result["chat_res"])
